package com.zw.chat.common.sse;

import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author by: zyl
 * @ClassName: SseSession
 * @Description:
 * @Date: 2024/11/12 15:25
 */
public class SseSession {

    /* Session维护 Map */
    private static final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();
    private static final Long MAX_TIME = 60_000L;

    /**
     * 获取SseEmitter
     *
     * @param clientId
     * @return
     */
    public static SseEmitter getEmitter(String clientId) {
        return sseEmitterMap.getOrDefault(clientId, new SseEmitter(MAX_TIME));
    }

    /**
     * 判断Session是否存在
     *
     * @param clientId
     * @return
     */
    public static boolean exists(String clientId) {
        return sseEmitterMap.containsKey(clientId);
    }

    /**
     * 增加Session
     *
     * @param clientId
     * @param emitter
     */
    public static void addEmitter(String clientId, SseEmitter emitter) {
        final SseEmitter sseEmitter = sseEmitterMap.get(clientId);
        if (sseEmitter != null) {
            sseEmitter.completeWithError(new RuntimeException("RepeatConnect(clientId:" + clientId + ")"));
        }
        sseEmitterMap.put(clientId, emitter);
    }

    /**
     * 移除Session
     *
     * @param clientId
     * @return
     */
    public static boolean rmEmitter(String clientId) {
        final SseEmitter sseEmitter = sseEmitterMap.remove(clientId);
        if (sseEmitter != null) {
            sseEmitter.complete();
            return true;
        }
        return false;
    }

    /**
     * SseEmitter onCompletion 后执行的逻辑
     *
     * @param clientId
     */
    public static void onCompletion(String clientId) {
        sseEmitterMap.remove(clientId);
    }

    /**
     * SseEmitter onError 或 onTimeout 后执行的逻辑
     *
     * @param clientId
     * @param e
     */
    public static void onError(String clientId, SseException e) {
        final SseEmitter sseEmitter = sseEmitterMap.get(clientId);
        if (sseEmitter != null) {
            sseEmitter.completeWithError(e);
        }
    }
}
